;;; Ported from Scheme 48 1.9.  See file COPYING for notices and license.
;;;
;;; Port Author: Andrew Whatson
;;;
;;; Original Authors: Richard Kelsey, Mike Sperber
;;;
;;;   scheme48-1.9.2/ps-compiler/prescheme/substitute.scm
;;;
;;; Substituting new variables for old in expressions.

(define-module (ps-compiler prescheme substitute)
  #:use-module (prescheme scheme48)
  #:use-module (prescheme s48-defrecord)
  #:use-module (prescheme bcomp node)
  #:use-module (prescheme bcomp binding)
  #:use-module (prescheme locations)
  #:use-module (ps-compiler node variable)
  #:use-module (ps-compiler prescheme eval)
  #:use-module (ps-compiler prescheme external-value)
  #:use-module (ps-compiler prescheme flatten)
  #:use-module ((ps-compiler prescheme form) #:select (note-variable-set!!))
  #:use-module (ps-compiler prescheme primitive)
  #:use-module (ps-compiler prescheme type)
  #:use-module (ps-compiler util strong)
  #:use-module (ps-compiler util util)
  #:export (substitute-in-expression
            note-variable-use!))

(define *free-exp-vars* #f)

(define (substitute-in-expression exp)
  (set! *free-exp-vars* '())
  (set! *letrec-datas* '())
  (let* ((exp (substitute-in-exp exp))
         (free *free-exp-vars*))
    (set! *free-exp-vars* '())
    (for-each (lambda (var)
                (set-variable-flag! var #f))
              free)
    (values exp free)))

(define global-marker (list 'global))

(define (note-binding-use! binding)
  (let ((var (binding-place binding)))
    (if (variable? var)
        (note-variable-use! var))))

(define (note-variable-use! var)
  (cond ((not (eq? (variable-flag var) global-marker))
         (set! *free-exp-vars* (cons var *free-exp-vars*))
         (set-variable-flag! var global-marker))))

;; Main dispatch

(define (substitute-in-exp node)
  ((operator-table-ref substitutions (node-operator-id node))
   node))

;; Particular operators

(define substitutions
  (make-operator-table
   (lambda (node)
     (error "no substitution for node ~S" node))))

(define (default-substitution node)
  (make-similar-node node
                     (cons (car (node-form node))
                           (map substitute-in-exp (cdr (node-form node))))))

(define (define-substitution name proc)
  (operator-define! substitutions name #f proc))

(define-substitution 'literal identity)
(define-substitution 'quote identity)
(define-substitution 'unspecific identity)

(define-substitution 'real-external
  (lambda (node)
    (let* ((exp (node-form node))
           (type (expand-type-spec (cadr (node-form (caddr exp))))))
      (make-literal-node (make-external-value (node-form (cadr exp))
                                              type)))))

(define op/literal (get-operator 'literal))

(define (make-literal-node x)
  (make-node op/literal x))

;; We copy the names because the same node may occur in multiple places
;; in the tree.

(define-substitution 'lambda
  (lambda (node)
    (let* ((new-names (copy-names (cadr (node-form node))))
           (body (substitute-in-exp (caddr (node-form node)))))
      (make-similar-node node
                         (list (car (node-form node))
                               new-names
                               body)))))

(define (copy-names names)
  (map (lambda (name)
         (let ((new (make-similar-node name (node-form name))))
           (node-set! name 'substitute new)
           new))
       names))

(define-substitution 'name
  (lambda (node)
    (substitute-name-node node #f)))

(define (substitute-name-node node call?)
  (let ((node (name-node-substitute node)))
    (let ((binding (node-ref node 'binding)))
      (cond ((not binding)
             (note-name-use! node)
             node)
            ((not (binding? binding))
             (bug "unbound variable ~S" (node-form node)))
            ((primitive? (binding-static binding))
             (make-primitive-node (binding-static binding) call?))
            ((location? (binding-place binding))
             (let ((value (contents (binding-place binding))))
               (if (constant? value)
                   (make-literal-node value)
                   (identity
                   (bug "name ~S has non-constant location ~S" node value)))))
            (else
             (note-binding-use! binding)
             node)))))

(define (name-node-substitute node)
  (let loop ((node node) (first? #t))
    (cond ((node-ref node 'substitute)
           => (lambda (node)
                (loop node #f)))
          ((and first? (not (node-ref node 'binding)))
           (user-error "unbound variable ~S" (node-form node)))
          (else
           node))))

(define-substitution 'set!
  (lambda (node)
    (let* ((exp (node-form node))
           (name (substitute-name-node (cadr exp) #f))
           (binding (node-ref name 'binding)))
      (if (not (binding? binding))
          (user-error "SET! on local variable ~S" (node-form (cadr exp))))
      (note-variable-set!!
        (binding-place binding))
      (note-binding-use! binding)
      (make-similar-node node
                         (list (car exp)
                               name
                               (substitute-in-exp (caddr exp)))))))

(define-substitution 'call
  (lambda (node)
    (let ((proc (car (node-form node)))
          (args (cdr (node-form node))))
      (make-similar-node node
                         (cons (if (name-node? proc)
                                   (substitute-name-node proc #t)
                                   (substitute-in-exp proc))
                               (map substitute-in-exp args))))))

;; Flush GOTO when it is used with a primitive.

(define-substitution 'goto
  (lambda (node)
    (let ((proc (cadr (node-form node)))
          (args (cddr (node-form node))))
      (if (and (name-node? proc)
               (bound-to-primitive? proc))
          (make-node (get-operator 'call)
                     (cons (substitute-name-node proc #t)
                           (map substitute-in-exp args)))
          (make-similar-node node
                             (cons 'goto
                                   (cons (if (name-node? proc)
                                             (substitute-name-node proc #t)
                                             (substitute-in-exp proc))
                                         (map substitute-in-exp args))))))))

(define name-node? (node-predicate 'name))

(define (bound-to-primitive? node)
  (let ((node (name-node-substitute node)))
    (let ((binding (node-ref node 'binding)))
      (and binding
           (primitive? (binding-static binding))))))

(define-substitution 'begin default-substitution)
(define-substitution 'if    default-substitution)

;; drop the loophole part

(define-substitution 'loophole
  (lambda (node)
    (substitute-in-exp (caddr (node-form node)))))

;;----------------------------------------------------------------
;; Breaking LETREC's down to improve type inference and make compilation
;; easier.

(define-substitution 'letrec
  (lambda (node)
    (let* ((exp (node-form node))
           (vars (map car (cadr exp)))
           (vals (map cadr (cadr exp))))
      (receive (names datas)
          (copy-letrec-names vars vals exp)
        (for-each (lambda (data value)
                    (expand-letrec-value data value datas exp))
                  datas
                  vals)
        (let ((sets (strongly-connected-components datas
                                                   letrec-data-uses
                                                   letrec-data-seen?
                                                   set-letrec-data-seen?!)))
          ;; so we don't keep track of which vars are referenced in the body
          (for-each (lambda (d)
                      (set-letrec-data-seen?! d #t))
                    datas)
          (do ((sets sets (cdr sets))
               (exp (substitute-in-exp (caddr exp))
                    (build-letrec (car sets) exp)))
              ((null? sets)
               (for-each (lambda (n)
                           (node-set! n 'letrec-data #f))
                         names)
               exp)))))))

(define-record-type letrec-data
  (name      ;; the name node for which this data exists
   marker    ;; a unique marker for this LETREC
   cell?     ;; variable is SET! or its value is not a (lambda ...).  This is
             ;; always #F until I can think of a reason to allow otherwise.
   )
  (value     ;; the expanded value of this variable
   uses      ;; a list of variables that VALUE uses
   seen?     ;; #T if this has been seen before during the current expansion
   ))

(define (copy-letrec-names names vals marker)
  (let ((names (map (lambda (name value)
                      (let ((new (make-similar-node name (node-form name)))
                            (cell? #f)) ;; we no longer allow SET! on LETREC vars.
                        (node-set! new 'letrec-data
                                   (letrec-data-maker new marker cell?))
                        (node-set! name 'substitute new)
                        new))
                    names
                    vals)))
    (values names (map (lambda (name) (node-ref name 'letrec-data)) names))))

(define lambda-node? (node-predicate 'lambda))

;; List of LETREC bound variables currently in scope.

(define *letrec-datas* '())

(define (note-name-use! name)
  (let ((data (node-ref name 'letrec-data)))
    (cond ((and data (not (letrec-data-seen? data)))
           (set-letrec-data-seen?! data #t)
           (set! *letrec-datas* (cons data *letrec-datas*))))))

;; Expand VALUE and determine which of DATAS it uses.

(define (expand-letrec-value data value datas mark)
  (let ((old-letrec-vars *letrec-datas*))
    (set! *letrec-datas* '())
    (for-each (lambda (d) (set-letrec-data-seen?! d #f)) datas)
    (set-letrec-data-value! data (substitute-in-exp value))
    (receive (ours others)
        (partition-list (lambda (data)
                          (eq? (letrec-data-marker data) mark))
                        *letrec-datas*)
      (set! *letrec-datas* (append others old-letrec-vars))
      (set-letrec-data-uses! data ours))))

;; If there is only one variable and its value doesn't reference it, then
;; use a LET instead of a LETREC.  Variables whose value is either set! or
;; not a lambda have explicit cells introduced.

(define (build-letrec datas body)
  (if (and (null? (cdr datas))
           (not (memq? (car datas)
                       (letrec-data-uses (car datas)))))
      (make-let-node (map letrec-data-name datas)
                     (map letrec-data-value datas)
                     body)
      (receive (cells normal)
          (partition-list letrec-data-cell? datas)
        (make-let-node (map letrec-data-name cells)
                       (map (lambda (ignore) (unspecific-node))
                            cells)
                       (make-letrec-node (map letrec-data-name normal)
                                         (map letrec-data-value normal)
                                         (make-begin-node
                                          (append (map letrec-data->set! cells)
                                                  (list body))))))))

(define op/unspecific (get-operator 'unspecific))
(define op/set! (get-operator 'set!))

(define (unspecific-node)
  (make-node op/unspecific '()))

(define (letrec-data->set! data)
  (make-node op/set!
             (list 'set!
                   (letrec-data-name data)
                   (letrec-data-value data))))

(define (make-let-node names values body)
  (if (null? names)
      body
      (make-node op/call
                 (cons (make-node op/lambda
                                  (list 'lambda names body))
                       values))))

(define (make-letrec-node names values body)
  (if (null? names)
      body
      (make-node op/letrec
                 (list 'letrec
                       (map list names values)
                       body))))

(define (make-begin-node nodes)
  (if (null? (cdr nodes))
      (car nodes)
      (make-node op/begin (cons 'begin nodes))))

(define op/call   (get-operator 'call))
(define op/lambda (get-operator 'lambda))
(define op/letrec (get-operator 'letrec))
(define op/begin  (get-operator 'begin))

;;----------------------------------------------------------------
;; A version of MAKE-SIMILAR-NODE that actually makes a new node.
;; I wish this could keep the old node's list of properties.

;; The version from nodes loops.

(define (make-similar-node node form)
  (make-node (node-operator node) form))
